import os, sys
import time
from astropy.io import fits
import numpy as np
from scipy.stats import sigmaclip
from astropy.table import Table


def remove_overscan(imgdata):
    print('> doing overscan correction ...')
    img = np.zeros((9232,9216), dtype=float)

    for i in range(4):
        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+27):((i+1)*1250-71)]
        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+27):((i+1)*1250-71)] 
        oc1 = np.mean(imgdata[0:4616,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float) 
        oc2 = np.mean(imgdata[4784:,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float)
#        print('oc1.shape: {}'.format(oc1.shape))
        for j in range(4616):
            img[j,i*1152:(i+1)*1152] = img[j,i*1152:(i+1)*1152] - oc1[j]
            img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152]- oc2[j]
    for i in range(4,8):
        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+71):((i+1)*1250-27)]
        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+71):((i+1)*1250-27)]
        oc1 = np.mean(imgdata[0:4616,(i*1250):(i*1250+71)], axis=1).astype(float)
        oc2 = np.mean(imgdata[4784:,(i*1250):(i*1250+71)], axis=1).astype(float)
        for j in range(4616):
            img[j, i*1152:(i+1)*1152] = img[j, i*1152:(i+1)*1152] - oc1[j]
            img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152] - oc2[j]

    print('> overscan correction done!')
    return img

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print('用法: python remove_overscan.py 输入FITS文件 [输出目录]')
        sys.exit(1)
    in_fits = sys.argv[1]
    if not os.path.isfile(in_fits):
        print(f'输入文件不存在: {in_fits}')
        sys.exit(1)
    if len(sys.argv) > 2:
        out_dir = sys.argv[2]
    else:
        out_dir = os.path.dirname(in_fits)
    os.makedirs(out_dir, exist_ok=True)
    base = os.path.splitext(os.path.basename(in_fits))[0]
    out_fits = os.path.join(out_dir, f'{base}_overscan_corrected.fits')

    with fits.open(in_fits) as hdul:
        imgdata = hdul[0].data
        header = hdul[0].header
        img_corr = remove_overscan(imgdata)
        fits.writeto(out_fits, img_corr, header, overwrite=True)
        print(f'> 修正后文件已保存: {out_fits}')